"""
图形识别核心算法
"""
import cv2
import numpy as np
from config.settings import SHAPE_DETECTION_CONFIG

class ShapeDetector:
    def __init__(self):
        self.min_contour_area = SHAPE_DETECTION_CONFIG["min_contour_area"]
        self.approx_epsilon_factor = SHAPE_DETECTION_CONFIG["approx_epsilon_factor"]
        self.blur_kernel_size = SHAPE_DETECTION_CONFIG["blur_kernel_size"]
    
    def detect_shapes(self, image):
        """识别图像中的图形"""
        # 预处理图像
        processed_image = self.preprocess_image(image)
        
        # 查找轮廓
        contours, _ = cv2.findContours(processed_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 创建结果图像
        # 修复：确保结果图像为彩色图像，而不是二值图像
        if len(image.shape) == 2:  # 如果是灰度图像
            result_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        else:
            result_image = image.copy()
        shapes = []
        
        # 获取图像尺寸
        height, width = image.shape[:2]
        image_area = height * width
        
        for contour in contours:
            # 过滤太小的轮廓
            area = cv2.contourArea(contour)
            if area < self.min_contour_area:
                continue
                
            # 过滤太大的轮廓（几乎占据整个图像）
            if area > 0.9 * image_area:
                continue
            
            # 识别图形
            shape_info = self.classify_shape(contour)
            if shape_info:
                shapes.append(shape_info)
                
                # 在图像上绘制结果
                self.draw_shape_info(result_image, contour, shape_info)
        
        return result_image, shapes
    
    def preprocess_image(self, image):
        """图像预处理"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(gray, (self.blur_kernel_size, self.blur_kernel_size), 0)
        
        # 自适应阈值二值化
        # 添加两种阈值方法，以适应不同大小的图形
        threshold = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 检查是否大部分区域都被标记为前景（可能是因为图形太大）
        # 如果超过90%的区域是白色（前景），则使用相反的阈值方法
        white_pixels = cv2.countNonZero(threshold)
        total_pixels = threshold.shape[0] * threshold.shape[1]
        if white_pixels > 0.9 * total_pixels:
            # 使用正向阈值二值化
            threshold = cv2.adaptiveThreshold(
                blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
            )
        
        return threshold
    
    def classify_shape(self, contour):
        """分类识别图形"""
        # 计算轮廓的周长
        perimeter = cv2.arcLength(contour, True)
        
        # 近似轮廓
        approx = cv2.approxPolyDP(contour, self.approx_epsilon_factor * perimeter, True)
        
        # 获取轮廓的边界矩形
        x, y, w, h = cv2.boundingRect(contour)
        
        # 计算长宽比
        aspect_ratio = float(w) / h
        
        # 计算面积
        area = cv2.contourArea(contour)
        
        # 计算圆形度
        circularity = 4 * np.pi * area / (perimeter * perimeter)
        
        # 根据特征分类
        shape_name = "unknown"
        confidence = 0.0
        
        # 判断顶点数量
        vertices = len(approx)
        
        if vertices == 3:
            # 三角形
            shape_name = "triangle"
            confidence = self.calculate_triangle_confidence(contour, approx)
            
        elif vertices == 4:
            # 四边形
            # 放宽正方形的判断条件，从0.95-1.05改为0.9-1.1
            if 0.9 <= aspect_ratio <= 1.1:
                shape_name = "square"
                confidence = self.calculate_square_confidence(contour, approx, aspect_ratio)
            else:
                shape_name = "rectangle"
                confidence = self.calculate_rectangle_confidence(contour, approx)
                
        elif vertices > 5 and circularity > 0.4:
            # 圆形（多边形近似时顶点很多且圆形度高）
            shape_name = "circle"
            confidence = self.calculate_circle_confidence(circularity, aspect_ratio)
        
        # 添加对近似顶点为3或4但实际是正方形的情况处理
        if confidence <= 0.5 and vertices >= 4:
            # 即使近似后顶点数大于4，也可能是一个正方形
            if 0.9 <= aspect_ratio <= 1.1:
                shape_name = "square"
                # 重新计算置信度，使用更宽松的规则
                confidence = self.calculate_square_confidence_loose(contour, aspect_ratio)
            elif aspect_ratio < 0.9 or aspect_ratio > 1.1:
                shape_name = "rectangle"
                confidence = 0.6  # 给一个基础置信度
        
        if confidence > 0.5:  # 只返回置信度较高的结果
            return {
                "name": shape_name,
                "confidence": confidence,
                "vertices": vertices,
                "area": area,
                "perimeter": perimeter,
                "aspect_ratio": aspect_ratio,
                "circularity": circularity
            }
        
        return None
    
    def calculate_triangle_confidence(self, contour, approx):
        """计算三角形识别的置信度"""
        # 基础置信度
        confidence = 0.8
        
        # 检查角度是否合理
        angles = self.calculate_angles(approx)
        if angles:
            angle_sum = sum(angles)
            if 170 < angle_sum < 190:  # 三角形内角和应该接近180度
                confidence += 0.1
            else:
                confidence -= 0.2
        
        return min(1.0, max(0.0, confidence))
    
    def calculate_square_confidence(self, contour, approx, aspect_ratio):
        """计算正方形识别的置信度"""
        confidence = 0.8
        
        # 长宽比越接近1，置信度越高
        ratio_score = 1 - abs(1 - aspect_ratio)
        confidence += ratio_score * 0.2
        
        # 检查角度是否接近90度
        angles = self.calculate_angles(approx)
        if angles:
            right_angle_count = sum(1 for angle in angles if 75 < angle < 105)
            confidence += (right_angle_count / 4) * 0.1
        
        return min(1.0, max(0.0, confidence))
    
    def calculate_square_confidence_loose(self, contour, aspect_ratio):
        """计算正方形识别的宽松置信度"""
        confidence = 0.7
        
        # 长宽比越接近1，置信度越高，但使用更宽松的标准
        ratio_score = 1 - abs(1 - aspect_ratio) * 1.5
        confidence += max(0, ratio_score * 0.3)
        
        # 面积较大的图形更可能是正方形
        area = cv2.contourArea(contour)
        if area > 2000:  # 面积较大时增加置信度
            confidence += 0.1
            
        return min(1.0, max(0.0, confidence))
    
    def calculate_rectangle_confidence(self, contour, approx):
        """计算长方形识别的置信度"""
        confidence = 0.7
        
        # 检查对边是否相等
        if len(approx) == 4:
            sides = self.calculate_side_lengths(approx)
            if len(sides) == 4:
                # 对边应该相等
                if abs(sides[0] - sides[2]) < 10 and abs(sides[1] - sides[3]) < 10:
                    confidence += 0.2
        
        return min(1.0, max(0.0, confidence))
    
    def calculate_circle_confidence(self, circularity, aspect_ratio):
        """计算圆形识别的置信度"""
        confidence = circularity  # 圆形度本身就是很好的置信度指标
        
        # 长宽比应该接近1
        if 0.8 <= aspect_ratio <= 1.2:
            confidence += 0.1
        
        return min(1.0, max(0.0, confidence))
    
    def calculate_angles(self, approx):
        """计算多边形的内角"""
        if len(approx) < 3:
            return []
        
        angles = []
        n = len(approx)
        
        for i in range(n):
            p1 = approx[i][0]
            p2 = approx[(i + 1) % n][0]
            p3 = approx[(i + 2) % n][0]
            
            # 计算向量
            v1 = p1 - p2
            v2 = p3 - p2
            
            # 计算角度
            cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
            cos_angle = np.clip(cos_angle, -1, 1)  # 避免数值错误
            angle = np.arccos(cos_angle) * 180 / np.pi
            angles.append(angle)
        
        return angles
    
    def calculate_side_lengths(self, approx):
        """计算多边形的边长"""
        if len(approx) < 3:
            return []
        
        sides = []
        n = len(approx)
        
        for i in range(n):
            p1 = approx[i][0]
            p2 = approx[(i + 1) % n][0]
            length = np.linalg.norm(p1 - p2)
            sides.append(length)
        
        return sides
    
    def draw_shape_info(self, image, contour, shape_info):
        """在图像上绘制图形信息"""
        # 绘制轮廓
        cv2.drawContours(image, [contour], -1, (0, 255, 0), 2)
        
        # 获取轮廓的中心点
        M = cv2.moments(contour)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
        else:
            cx, cy = 0, 0
        
        # 绘制标签
        label = f"{shape_info['name']} ({shape_info['confidence']:.2f})"
        
        # 计算文本大小
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.6
        thickness = 2
        (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
        
        # 绘制背景矩形
        cv2.rectangle(image, 
                     (cx - text_width // 2 - 5, cy - text_height - 5),
                     (cx + text_width // 2 + 5, cy + 5),
                     (255, 255, 255), -1)
        
        # 绘制文本
        cv2.putText(image, label, (cx - text_width // 2, cy), 
                   font, font_scale, (0, 0, 0), thickness)
        
        # 绘制中心点
        cv2.circle(image, (cx, cy), 3, (255, 0, 0), -1)
    
    def get_shape_by_name(self, image, target_shape):
        """获取指定名称的图形"""
        _, shapes = self.detect_shapes(image)
        
        for shape in shapes:
            if shape["name"].lower() == target_shape.lower():
                return shape
        
        return None